import logging
import socket
import sys
import torch
import numpy as np
from PIL import Image
import io
import pandas as pd

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 10, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(10, 20, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(320, 50),
            torch.nn.Linear(50, 10),
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv1(x)  # 一层卷积层,一层池化层,一层激活层(图是先卷积后激活再池化，差别不大)
        x = self.conv2(x)  # 再来一次
        x = x.view(batch_size, -1)  # flatten 变成全连接网络需要的输入 (batch, 20,4,4) ==> (batch,320), -1 此处自动算出的是320
        x = self.fc(x)
        return x  # 最后输出的是维度为10的，也就是（对应数学符号的0~9）


def process_image(image_data):
    img = Image.open(io.BytesIO(image_data))
    img_array = np.array(img)
    model_input = torch.from_numpy(img_array).float().to(device)
    model = torch.load('./model_Mnist_optimal.pth', map_location=torch.device(device))
    model_input = model_input.unsqueeze(0)
    print(model_input.size())
    pre_res = model(model_input)
    pre_res = torch.argmax(pre_res, dim=1)
    return pre_res.item()


if __name__ == "__main__":


    hostname = socket.gethostname()
    print("系统主机名:", hostname)

    df = pd.read_parquet('train.parquet', engine='pyarrow')
    img = df.iloc[0]['image']['bytes']
    label = df.iloc[0]['label']
    print(label)
    print(process_image(img))
